#频谱中心内化后提取低中高三频段然后分别映射回去查看
# -*- coding: utf-8 -*-
import numpy as np
import cv2
import matplotlib.pyplot as plt

def frequency_masks(shape, low_ratio=0.1, high_ratio=0.3):
    H, W = shape
    center_u, center_v = H // 2, W // 2
    U, V = np.ogrid[:H, :W]
    dist = np.sqrt((U - center_u)**2 + (V - center_v)**2)
    max_dist = dist.max()

    low_mask = dist <= low_ratio * max_dist
    mid_mask = (dist > low_ratio * max_dist) & (dist <= high_ratio * max_dist)
    high_mask = dist > high_ratio * max_dist
    return low_mask, mid_mask, high_mask

def process_image_frequency_bands(img_path):
    # 读取灰度图
    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
    H, W = img.shape

    # 计算二维DFT并中心化频谱
    F = np.fft.fft2(img)
    F_shifted = np.fft.fftshift(F)

    # 获取低中高频掩膜
    low_mask, mid_mask, high_mask = frequency_masks((H, W))

    # 提取不同频段
    F_low = F_shifted * low_mask
    F_mid = F_shifted * mid_mask
    F_high = F_shifted * high_mask

    # 逆变换恢复图像
    img_low = np.fft.ifft2(np.fft.ifftshift(F_low)).real
    img_mid = np.fft.ifft2(np.fft.ifftshift(F_mid)).real
    img_high = np.fft.ifft2(np.fft.ifftshift(F_high)).real

    # 显示结果
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 黑体字体，Windows系统常见
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
    plt.figure(figsize=(12,8))
    plt.subplot(2,3,1); plt.imshow(img, cmap='gray'); plt.title('原图'); plt.axis('off')
    plt.subplot(2,3,2); plt.imshow(np.log(np.abs(F_shifted)+1), cmap='gray'); plt.title('频谱（对数）'); plt.axis('off')
    plt.subplot(2,3,4); plt.imshow(img_low, cmap='gray'); plt.title('低频部分'); plt.axis('off')
    plt.subplot(2,3,5); plt.imshow(img_mid, cmap='gray'); plt.title('中频部分'); plt.axis('off')
    plt.subplot(2,3,6); plt.imshow(img_high, cmap='gray'); plt.title('高频部分'); plt.axis('off')

    plt.tight_layout()
    plt.show()

# 这里替换成你本地的图片路径
process_image_frequency_bands('D:\\2D-DFT\image\\testhaze4.png')
